# dataset settings
from mmcv.transforms import LoadImageFromFile, RandomResize
from mmengine.dataset import DefaultSampler

from mmdet.datasets import AspectRatioBatchSampler
from mmdet.datasets.transforms import RandomFlip, RandomCrop, PackDetInputs, Resize
from mmdet.evaluation import CocoPanopticMetric

from seg.datasets.ade_ov import ADEPanopticOVDataset
from seg.datasets.pipelines.loading import LoadPanopticAnnotationsHB

data_root = 'data/ade/'
backend_args = None
image_size = (1280, 736)

train_pipeline = [
    dict(
        type=LoadImageFromFile,
        to_float32=True,
        backend_args=backend_args),
    dict(
        type=LoadPanopticAnnotationsHB,
        with_bbox=True,
        with_mask=True,
        with_seg=True,
        backend_args=backend_args),
    dict(type=RandomFlip, prob=0.5),
    dict(
        type=RandomResize,
        resize_type=Resize,
        scale=image_size,
        ratio_range=(0.1, 2.0),
        keep_ratio=True,
    ),
    dict(
        type=RandomCrop,
        crop_size=image_size,
        crop_type='absolute',
        recompute_bbox=True,
        allow_negative_crop=True),
    dict(type=PackDetInputs)
]
train_dataloader = dict(
    batch_size=2,
    num_workers=2,
    persistent_workers=True,
    sampler=dict(type=DefaultSampler, shuffle=True),
    batch_sampler=dict(type=AspectRatioBatchSampler),
    dataset=dict(
        type=ADEPanopticOVDataset,
        data_root=data_root,
        ann_file='ADEChallengeData2016/ade20k_panoptic_train.json',
        data_prefix=dict(img='ADEChallengeData2016/images/training/',
                         seg='ADEChallengeData2016/ade20k_panoptic_train/'),
        filter_cfg=dict(filter_empty_gt=True, min_size=32),
        pipeline=train_pipeline,
        backend_args=backend_args
    )
)

test_pipeline = [
    dict(type=LoadImageFromFile, backend_args=backend_args),
    dict(type=Resize, scale=(2560, 640), keep_ratio=True),
    dict(type=LoadPanopticAnnotationsHB, backend_args=backend_args),
    dict(
        type=PackDetInputs,
        meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', 'scale_factor')
    )
]
val_dataloader = dict(
    batch_size=2,
    num_workers=2,
    persistent_workers=True,
    drop_last=False,
    sampler=dict(type=DefaultSampler, shuffle=False),
    dataset=dict(
        type=ADEPanopticOVDataset,
        data_root=data_root,
        ann_file='ADEChallengeData2016/ade20k_panoptic_val.json',
        data_prefix=dict(img='ADEChallengeData2016/images/validation/',
                         seg='ADEChallengeData2016/ade20k_panoptic_val/'),
        test_mode=True,
        pipeline=test_pipeline,
        backend_args=backend_args
    )
)
test_dataloader = val_dataloader

val_evaluator = dict(
    type=CocoPanopticMetric,
    # ann_file=data_root + 'ADEChallengeData2016/ade20k_panoptic_val.json',
    seg_prefix=data_root + 'ADEChallengeData2016/ade20k_panoptic_val/',
    backend_args=backend_args
)
test_evaluator = val_evaluator
